#
import sys
#import matplotlib
#matplotlib.use('Agg')  # 使用非交互式后端，避免 DISPLAY 问题
import matplotlib.pyplot as plt
import numpy
import pysam


def parse_bam(bam_path):
    errors = []
    qualities = []
    lengths = []
    read_count = 0
    samfile = pysam.AlignmentFile(bam_path, "rb")
    for read in samfile.fetch():
        read_count += 1
        if read.is_duplicate or read.is_secondary or read.is_supplementary or read.is_unmapped or not read.has_tag("NM"):
            continue
        query_alignment_length = read.query_alignment_length
        nm = read.get_tag("NM")
        mapq = read.mapping_quality
        query_length = read.query_length
        error_rate = nm / query_alignment_length * 100
        qualities.append(mapq)
        lengths.append(query_length)
        errors.append(error_rate)
    avg_length = sum(lengths) / len(lengths) if lengths else 0
    mean_errors = sum(errors) / len(errors) if errors else 0  # 计算错误率平均值
    return errors, qualities, lengths, read_count, avg_length, mean_errors


def plot_error(errors, output_path="error_rate.png"):
    plt.figure()
    plt.hist(errors, range=(0, 40), bins=200)
    plt.xlabel("error rate (%)")
    plt.ylabel("number of reads")
    plt.savefig(output_path)
    plt.close()


def n50(lengths):
    all_len = sorted(lengths, reverse=True)
    csum = numpy.cumsum(all_len)
    print("N: %d" % int(sum(lengths)))
    n2 = int(sum(lengths) / 2)
    csumn2 = min(csum[csum >= n2])
    ind = numpy.where(csum == csumn2)
    n50 = all_len[ind[0][0]]
    print("N50: %s" % n50)


if __name__ == '__main__':
    if len(sys.argv) < 2:
        print("Usage: python script.py <bam_file> [output_image_path]")
        sys.exit(1)
    
    bam_path = sys.argv[1]
    # 如果提供了输出路径，使用它；否则使用默认值
    output_image_path = sys.argv[2] if len(sys.argv) > 2 else "error_rate.png"
    
    errors, qualities, lengths, read_count, avg_length, mean_errors = parse_bam(bam_path)
    print(f"NO. of reads: {read_count}")  # 恢复之前的输出格式
    print(f"mean errors: {mean_errors}")
    print(f"Total reads processed: {read_count}")
    print(f"Average read length: {avg_length:.2f}")
    plot_error(errors, output_image_path)
    n50(lengths)
